import pandas as pd
import numpy as np
from scipy.stats import pearsonr

# 定义计算指标的函数
def calculate_metrics(ground_precipitation, other_precipitation):
    """
    计算平均偏差、平均绝对误差、均方根误差、平均相对误差、相关系数。

    参数:
    ground_precipitation (pd.Series): 地面 1 小时降水数据
    other_precipitation (pd.Series): 其他降水数据

    返回:
    dict: 包含五个评估指标的字典
    """
    # 处理地面站降水为 0 的情况，加上一个很小的数
    ground_precipitation = ground_precipitation.replace(0, 1e-10)

    if len(ground_precipitation) < 2 or len(other_precipitation) < 2:
        # 数据长度不足 2，用 np.nan 填充指标
        return {
            '平均偏差': np.nan,
            '平均绝对误差': np.nan,
            '均方根误差': np.nan,
            '平均相对误差(%)': np.nan,
            '相关系数': np.nan
        }
    
    bias = np.mean(other_precipitation - ground_precipitation)
    mae = np.mean(np.abs(other_precipitation - ground_precipitation))
    rmse = np.sqrt(np.mean((other_precipitation - ground_precipitation) ** 2))
    mre = np.mean(np.abs((other_precipitation - ground_precipitation) / ground_precipitation)) * 100
    corr, _ = pearsonr(ground_precipitation, other_precipitation)
    
    return {
        '平均偏差': bias,
        '平均绝对误差': mae,
        '均方根误差': rmse,
        '平均相对误差(%)': mre,
        '相关系数': corr
    }

def read_and_merge_data(file_path, columns_to_analyze):
    """
    读取 Excel 文件并合并数据，同时过滤缺失值。

    参数:
    file_path (str): Excel 文件路径
    columns_to_analyze (list): 需要分析的列名列表

    返回:
    dict: 合并后的数据字典
    """
    excel_file = pd.ExcelFile(file_path)
    merged_data = {col: {'ground': [], 'other': []} for col in columns_to_analyze}
    total_original_count = 0
    total_filtered_count = 0

    for sheet_name in excel_file.sheet_names:
        df = excel_file.parse(sheet_name)
        ground_precipitation = df["地面1小时降水"]
        for column in columns_to_analyze:
            if column in df.columns:
                other_precipitation = df[column]
                original_count = len(ground_precipitation)
                total_original_count += original_count
                # 在合并前过滤缺失值
                valid_mask = ~ground_precipitation.isnull() & ~other_precipitation.isnull()
                merged_ground = ground_precipitation[valid_mask]
                merged_other = other_precipitation[valid_mask]
                filtered_count = len(merged_ground)
                total_filtered_count += filtered_count
                merged_data[column]['ground'].append(merged_ground)
                merged_data[column]['other'].append(merged_other)
                # print(f"工作表 {sheet_name}，列 {column}，原始数据样本数: {original_count}, 过滤后数据样本数: {filtered_count}")
    print(f"所有工作表合并后，原始数据总样本数: {total_original_count}, 过滤后数据总样本数: {total_filtered_count}")
    return merged_data

def main():
    # 读取 Excel 文件
    # file_path = "statistic/2024-07-16日00至17日23时降水数据下载.xlsx"
    file_path = "statistic/2025-05-08日18时至21时降水数据下载.xlsx"

    # 定义需要分析的列名
    columns_to_analyze = [
        # "雷达降水估测",
        "雷达拼图3.0降水评估",
        # "CMPAS产品-三源融合快速1小时降水", 
        # "CMPAS分省逐小时降水实时",
        # "CMPAS分省逐小时降水近实时产品"
    ]

    # 读取并合并数据
    merged_data = read_and_merge_data(file_path, columns_to_analyze)

    # 计算合并后的数据指标
    results = {}
    for column in columns_to_analyze:
        merged_ground = pd.concat(merged_data[column]['ground'])
        merged_other = pd.concat(merged_data[column]['other'])
        metrics = calculate_metrics(merged_ground, merged_other)
        results[column] = metrics

    # 将过滤后的数据输出到 Excel 文件
    filtered_data_file = "分产品输出0508.xlsx"
    # 修改引擎为 openpyxl
    with pd.ExcelWriter(f"statistic/{filtered_data_file}", engine='openpyxl') as writer:
        for column in columns_to_analyze:
            merged_ground = pd.concat(merged_data[column]['ground'])
            merged_other = pd.concat(merged_data[column]['other'])
            df_filtered = pd.DataFrame({
                '地面1小时降水': merged_ground,
                column: merged_other
            })
            df_filtered.to_excel(writer, sheet_name=column, index=False)
    print(f"过滤后的数据已写入 {filtered_data_file}")

    # 将结果转换为 DataFrame 并写入 Excel
    output_data = []
    for column, metrics in results.items():
        row = {'产品名称': column, **metrics}
        output_data.append(row)

    output_df = pd.DataFrame(output_data)
    output_file_path = "降水评估结果_合并版0508.xlsx"
    output_df.to_excel(f"statistic/{output_file_path}", index=False)
    print(f"评估结果已写入 {output_file_path}")

if __name__ == "__main__":
    main()